#include "linkfabric/MemLink.hpp"
#include <fcntl.h>
#include <sys/mman.h>
#include <unistd.h>
#include <cstring>
#include "linkfabric/Futex.hpp"

using namespace linkfabric;

bool MemLink::init(const std::string& devicePath) {
    fd_ = open(devicePath.c_str(), O_RDWR);
    if (fd_ < 0) return false;
    shm_ = reinterpret_cast<SharedMemory*>(
        mmap(nullptr, sizeof(SharedMemory),
             PROT_READ | PROT_WRITE, MAP_SHARED, fd_, 0));
    return shm_ != MAP_FAILED;
}

int MemLink::allocChannel(ChannelPermission perm) {
    for (int i = 0; i < MAX_CHANNELS; ++i) {
        uint32_t mask = 1u << i;
        if (!(shm_->channel_bitmap.load() & mask)) {
            shm_->channel_bitmap |= mask;
            shm_->channels[i].permission = perm;
            shm_->channels[i].futex_flag.store(0);
            return i;
        }
    }
    return -1;
}

void MemLink::freeChannel(int channelId) {
    uint32_t mask = ~(1u << channelId);
    shm_->channel_bitmap &= mask;
    shm_->channels[channelId].permission = ChannelPermission::NONE;
}

bool MemLink::write(int channelId, const std::vector<uint8_t>& data) {
    Channel& ch = shm_->channels[channelId];
    if ((ch.permission & ChannelPermission::WRITE) == ChannelPermission::NONE) return false;
    if (data.size() > MAX_PAYLOAD) return false;
    if (ch.futex_flag.load()) return false;
    ch.length = static_cast<uint32_t>(data.size());
    memcpy(ch.payload, data.data(), data.size());
    ch.futex_flag.store(1);
    futex_wake(&ch.futex_flag);
    return true;
}

void MemLink::onData(int channelId, std::function<void(const std::vector<uint8_t>&)> cb) {
    callbacks_[channelId] = cb;
}

void MemLink::run() {
    while (true) {
        for (auto& [chId, cb] : callbacks_) {
            Channel& ch = shm_->channels[chId];
            if ((ch.permission & ChannelPermission::READ) == ChannelPermission::NONE) continue;
            if (ch.futex_flag.load() == 0) {
                futex_wait(&ch.futex_flag, 0);
            }
            std::vector<uint8_t> buf(ch.payload, ch.payload + ch.length);
            cb(buf);
            ch.futex_flag.store(0);
        }
        usleep(1000);
    }
}

void MemLink::shutdown() {
    munmap(shm_, sizeof(SharedMemory));
    close(fd_);
}